"""
Function Calling Schema定义
定义所有工具函数的JSON Schema
"""

from typing import List, Dict, Any


def get_calculator_schema() -> Dict[str, Any]:
    """获取基本计算器的Schema"""
    return {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "基本数学计算器，支持加减乘除、幂运算、取模等基本运算",
            "parameters": {
                "type": "object",
                "properties": {
                    "op": {
                        "type": "string",
                        "enum": ["add", "sub", "mul", "div", "pow", "mod"],
                        "description": "运算类型：add(加法), sub(减法), mul(乘法), div(除法), pow(幂运算), mod(取模)"
                    },
                    "a": {
                        "type": "number",
                        "description": "第一个操作数"
                    },
                    "b": {
                        "type": "number", 
                        "description": "第二个操作数"
                    }
                },
                "required": ["op", "a", "b"]
            }
        }
    }


def get_advanced_calculator_schema() -> Dict[str, Any]:
    """获取高级计算器的Schema"""
    return {
        "type": "function",
        "function": {
            "name": "advanced_calculator",
            "description": "高级数学计算器，支持三角函数、对数、开方、绝对值、四舍五入等高级运算",
            "parameters": {
                "type": "object",
                "properties": {
                    "operation": {
                        "type": "string",
                        "enum": ["sqrt", "sin", "cos", "tan", "log", "ln", "abs", "round"],
                        "description": "运算类型：sqrt(平方根), sin(正弦), cos(余弦), tan(正切), log(对数), ln(自然对数), abs(绝对值), round(四舍五入)"
                    },
                    "args": {
                        "type": "array",
                        "items": {
                            "type": "number"
                        },
                        "description": "运算参数列表，不同运算需要的参数个数不同"
                    }
                },
                "required": ["operation", "args"]
            }
        }
    }


def get_music_player_schemas() -> List[Dict[str, Any]]:
    """获取音乐播放器相关函数的Schema"""
    return [
        {
            "type": "function",
            "function": {
                "name": "music_list",
                "description": "获取可播放的音乐列表",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_play",
                "description": "播放指定的歌曲",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "music_name": {
                            "type": "string",
                            "description": "歌曲名称，可以带或不带扩展名"
                        }
                    },
                    "required": ["music_name"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_play_random",
                "description": "随机播放一首歌曲",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_stop",
                "description": "停止当前播放的音乐",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_volume_set",
                "description": "设置音乐播放音量（0.0-1.0）",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "volume": {
                            "type": "number",
                            "description": "音量大小，范围0.0（静音）到1.0（最大音量）",
                            "minimum": 0.0,
                            "maximum": 1.0
                        }
                    },
                    "required": ["volume"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_volume_up",
                "description": "增加音乐播放音量",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "step": {
                            "type": "number",
                            "description": "音量增加的步长，默认0.1",
                            "default": 0.1,
                            "minimum": 0.01,
                            "maximum": 1.0
                        }
                    }
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_volume_down",
                "description": "减少音乐播放音量",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "step": {
                            "type": "number",
                            "description": "音量减少的步长，默认0.1",
                            "default": 0.1,
                            "minimum": 0.01,
                            "maximum": 1.0
                        }
                    }
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_volume_set_smart",
                "description": "智能设置音乐音量（支持数字、百分比、中文描述）",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "value": {
                            "type": "string",
                            "description": "音量值：数字(如0.8)、百分比(如'50%')或中文描述(如'最大'、'一半'、'静音')"
                        }
                    },
                    "required": ["value"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "music_status",
                "description": "获取当前音乐播放状态",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        }
    ]


def get_all_schemas() -> List[Dict[str, Any]]:
    """获取所有工具函数的Schema"""
    schemas = []
    
    # 添加基本计算器
    schemas.append(get_calculator_schema())
    
    # 添加高级计算器
    schemas.append(get_advanced_calculator_schema())
    
    # 添加音乐播放器
    schemas.extend(get_music_player_schemas())
    
    return schemas


def get_gimbal_schemas() -> List[Dict[str, Any]]:
    """获取云台控制相关函数的Schema"""
    return [
        {
            "type": "function",
            "function": {
                "name": "gimbal_control",
                "description": "控制云台转动到指定的绝对角度位置",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "pan": {
                            "type": "number",
                            "description": "水平角度，范围0-270度（0度=最左，135度=中心，270度=最右）",
                            "minimum": 0,
                            "maximum": 270
                        },
                        "tilt": {
                            "type": "number",
                            "description": "垂直角度，范围0-270度（0度=最下，135度=中心，270度=最上）",
                            "minimum": 0,
                            "maximum": 270
                        }
                    },
                    "required": ["pan", "tilt"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "gimbal_move_relative",
                "description": "相对移动云台，在当前位置基础上增加指定的角度偏移",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "pan_delta": {
                            "type": "number",
                            "description": "水平角度变化量（正值向右，负值向左）",
                            "default": 0
                        },
                        "tilt_delta": {
                            "type": "number",
                            "description": "垂直角度变化量（正值向上，负值向下）",
                            "default": 0
                        }
                    },
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "gimbal_reset",
                "description": "重置云台到初始中心位置",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "gimbal_status",
                "description": "获取云台当前状态信息",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "gimbal_connect",
                "description": "连接云台控制器",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "gimbal_disconnect",
                "description": "断开云台控制器连接",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        }
    ]


def get_arm_control_schemas() -> List[Dict[str, Any]]:
    """获取机械臂控制相关函数的Schema"""
    return [
        {
            "type": "function",
            "function": {
                "name": "arm_catch_object",
                "description": "控制机械臂抓取物体",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "sth_position_x": {
                            "type": "number",
                            "description": "物体在图像中的x坐标"
                        },
                        "sth_position_y": {
                            "type": "number",
                            "description": "物体在图像中的y坐标"
                        },
                        "sth_position_z": {
                            "type": "number",
                            "description": "抓取高度，默认-50.0",
                            "default": -50.0
                        }
                    },
                    "required": ["sth_position_x", "sth_position_y"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_put_object",
                "description": "控制机械臂放置物体",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "sth_position_x": {
                            "type": "number",
                            "description": "放置位置的x坐标"
                        },
                        "sth_position_y": {
                            "type": "number",
                            "description": "放置位置的y坐标"
                        },
                        "sth_position_z": {
                            "type": "number",
                            "description": "放置高度，默认-50.0",
                            "default": -50.0
                        }
                    },
                    "required": ["sth_position_x", "sth_position_y"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_stack_put_object",
                "description": "控制机械臂堆叠放置物体",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "sth_position_x": {
                            "type": "number",
                            "description": "放置位置的x坐标"
                        },
                        "sth_position_y": {
                            "type": "number",
                            "description": "放置位置的y坐标"
                        },
                        "sth_position_z": {
                            "type": "number",
                            "description": "堆叠放置高度，默认30.0",
                            "default": 30.0
                        }
                    },
                    "required": ["sth_position_x", "sth_position_y"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_init_action",
                "description": "执行机械臂初始化动作",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_move_direction",
                "description": "控制机械臂向指定方向移动",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "direction": {
                            "type": "string",
                            "enum": ["forward", "back", "left", "right", "up", "down"],
                            "description": "移动方向"
                        },
                        "distance": {
                            "type": "number",
                            "description": "移动距离（厘米）"
                        }
                    },
                    "required": ["direction", "distance"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_gripper_control",
                "description": "控制机械臂夹爪开合",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "action": {
                            "type": "string",
                            "enum": ["open", "close"],
                            "description": "动作类型"
                        }
                    },
                    "required": ["action"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_status",
                "description": "获取机械臂当前状态信息",
                "parameters": {
                    "type": "object",
                    "properties": {}
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "arm_move_to_position",
                "description": "控制机械臂移动到指定的绝对坐标位置(x,y,z)，适用于\"移动到\"、\"到达位置\"等绝对位置指令",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "target_x": {
                            "type": "number",
                            "description": "目标X坐标 (毫米)，范围50-180"
                        },
                        "target_y": {
                            "type": "number", 
                            "description": "目标Y坐标 (毫米)，范围-80-80"
                        },
                        "target_z": {
                            "type": "number",
                            "description": "目标Z坐标 (毫米)，范围-50-140"
                        }
                    },
                    "required": ["target_x", "target_y", "target_z"]
                }
            }
        }
    ]


def get_vlm_schemas() -> List[Dict[str, Any]]:
    """获取视觉语言模型相关函数的Schema"""
    return [
        {
            "type": "function",
            "function": {
                "name": "vlm_analyze_image",
                "description": "分析图像内容，提供详细的图像描述和理解",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "image_base64": {
                            "type": "string",
                            "description": "图像的base64编码字符串，应该包含data:image/前缀"
                        },
                        "prompt": {
                            "type": "string",
                            "description": "分析提示词，指导AI如何分析图像",
                            "default": "请详细描述这张图片的内容"
                        }
                    },
                    "required": ["image_base64"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "vlm_detect_objects",
                "description": "检测图像中的物体并提供边界框坐标",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "image_base64": {
                            "type": "string",
                            "description": "图像的base64编码字符串，应该包含data:image/前缀"
                        },
                        "prompt": {
                            "type": "string",
                            "description": "检测提示词，指导AI如何检测物体",
                            "default": "请检测图像中的物体并提供边界框坐标"
                        }
                    },
                    "required": ["image_base64"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "vlm_answer_visual_question",
                "description": "回答关于图像内容的具体问题",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "image_base64": {
                            "type": "string",
                            "description": "图像的base64编码字符串，应该包含data:image/前缀"
                        },
                        "question": {
                            "type": "string",
                            "description": "关于图像的问题"
                        }
                    },
                    "required": ["image_base64", "question"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "vlm_read_text",
                "description": "识别并提取图像中的文字内容",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "image_base64": {
                            "type": "string",
                            "description": "图像的base64编码字符串，应该包含data:image/前缀"
                        }
                    },
                    "required": ["image_base64"]
                }
            }
        }
    ]


def get_enabled_schemas(config) -> List[Dict[str, Any]]:
    """根据配置获取启用的工具Schema"""
    schemas = []
    
    if not config.function_calling.enabled:
        return schemas
    
    # 检查计算器是否启用
    calculator_config = config.function_calling.tools.get('calculator')
    if calculator_config and calculator_config.enabled:
        schemas.append(get_calculator_schema())
        schemas.append(get_advanced_calculator_schema())
    
    # 检查音乐播放器是否启用
    music_player_config = config.function_calling.tools.get('music_player')
    if music_player_config and music_player_config.enabled:
        schemas.extend(get_music_player_schemas())
    
    # 检查云台控制是否启用
    gimbal_config = config.function_calling.tools.get('gimbal')
    if gimbal_config and gimbal_config.enabled:
        schemas.extend(get_gimbal_schemas())
    
    # 检查机械臂控制是否启用
    arm_control_config = config.function_calling.tools.get('arm_control')
    if arm_control_config and arm_control_config.enabled:
        schemas.extend(get_arm_control_schemas())
    
    # 检查VLM功能是否启用
    if config.vlm.enabled:
        schemas.extend(get_vlm_schemas())
    
    return schemas
